from utils.preprocessing import *
from utils.train import train, trainlog
from models.fucknet import fuckSeg, fuckSeg2
from torch.nn import BCEWithLogitsLoss
from torch.optim import lr_scheduler,Adam,RMSprop
from utils.preprocessing import gen_dataloader
from models.GCN import *
from models.losses import CrossEntropyLoss2d
from models.icnet import icnet_portrait
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
img_root = '/media/hszc/data1/seg_data/'

resume = '/home/hszc/zhangchi/models/FigureSeg/icnet-385/bestweights-[0.863].pth'


data_set, data_loader = gen_dataloader(img_root, validation_split=0.1, train_bs=8, val_bs=24)
print len(data_set['train']), len(data_set['val'])

model = icnet_portrait(num_classes=2)


if resume:
    model.eval()
    logging.info('resuming finetune from %s' % resume)
    try:
        model.load_state_dict(torch.load(resume))
    except KeyError:
        model = torch.nn.DataParallel(model)
        model.load_state_dict(torch.load(resume))

import time
from torch.autograd import Variable



imgs, masks = data_set['train'][0]
imgs = imgs.unsqueeze(0)
imgs = Variable(imgs)
t0 = time.time()
print imgs.size()
for i in range(100):
    print i
    model(imgs)
t1 = time.time()
print t1-t0
